Skip to content

[converter] Fix negative axis in quantize/dequantize lowering - #24

Merged
jakesabathia2 merged 3 commits into
apple:mainfrom
eyupcanakman:fix/dequantize-negative-axis
Jul 8, 2026
Merged

[converter] Fix negative axis in quantize/dequantize lowering#24
jakesabathia2 merged 3 commits into
apple:mainfrom
eyupcanakman:fix/dequantize-negative-axis

Conversation

@eyupcanakman

Copy link
Copy Markdown
Contributor

Description:

  • The coreai::quantize / coreai::dequantize lowering normalized a negative axis as axis + rank - 1, off by one from the eager op, which resolves it as axis + rank. A per-channel axis=-1 then landed one dimension early. When the channel and a neighbor dim share a size there is no shape error, so the converted model picks the wrong channel silently.
  • Normalize a negative axis the same way the eager op does (axis + rank), so the converted program matches eager for every negative axis.

Testing

  • python unit test
  • Added per-channel negative-axis numerical tests for quantize and dequantize that fail on the old formula and pass with the fix.

The quantize/dequantize lowering normalized a negative axis as
axis + rank - 1, off by one from the eager op, which resolves it as
axis + rank. A per-channel axis=-1 was applied one dimension early, so
the converted model used the wrong channel; when the channel and a
neighbor dim share a size this is silent, with no shape error.

Normalize a negative axis as axis + rank, matching the eager op. Add
per-channel negative-axis numerical tests for both ops. They use a
(2, 4, 4) shape with equal middle and channel dims so a wrong axis
surfaces as a numerical mismatch rather than a reshape error.
Comment thread tests/ops/test_custom_ops.py Outdated
)

model = Model()
x = torch.randn(2, 4, 4)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we can use a different dim value at [1], so something like (2,3,4) for clarity on which axis is picked via the negative indexing.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added (2, 3, 4) as you suggested and kept (2, 4, 4) too. Equal dims are the case where a wrong axis is silent instead of a reshape error, so both stay covered.

Comment thread tests/ops/test_custom_ops.py Outdated
)

model = Model()
x = torch.randint(-128, 127, (2, 4, 4), dtype=torch.int8)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same comment as above

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here.

@jakesabathia2

Copy link
Copy Markdown
Contributor

Thank you for the fix @eyupcanakman , can you please rebase on the latest main and re-push?

@pkmandke pkmandke left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lgtm!
@yixingli-apple can you please take a look as well?

@jakesabathia2
jakesabathia2 merged commit 40312b5 into apple:main Jul 8, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants